'''
 * @Author: Benjay·Shaw
 * @Date: 2024-10-31 17:07:50
 * @LastEditors: Benjay·Shaw
 * @LastEditTime: 2024-10-31 22:15:47
 * @Description: VGG
'''
import paddle


class Vgg16(paddle.nn.Layer):

    def __init__(self):
        super(Vgg16, self).__init__()
        features = paddle.vision.models.vgg16(pretrained=True).features
        self.to_relu_1_2 = paddle.nn.Sequential()
        self.to_relu_2_2 = paddle.nn.Sequential()
        self.to_relu_3_3 = paddle.nn.Sequential()
        self.to_relu_4_3 = paddle.nn.Sequential()
        self.to_relu_5_3 = paddle.nn.Sequential()
        for x in range(4):
            self.to_relu_1_2.add_sublayer(name=str(x), sublayer=features[x])
        for x in range(4, 9):
            self.to_relu_2_2.add_sublayer(name=str(x), sublayer=features[x])
        for x in range(9, 16):
            self.to_relu_3_3.add_sublayer(name=str(x), sublayer=features[x])
        for x in range(16, 23):
            self.to_relu_4_3.add_sublayer(name=str(x), sublayer=features[x])
        for x in range(23, 30):
            self.to_relu_5_3.add_sublayer(name=str(x), sublayer=features[x])
        for param in self.parameters():
            param.stop_gradient = not False

    def forward(self, x):
        h = self.to_relu_1_2(x)
        h_relu_1_2 = h
        h = self.to_relu_2_2(h)
        h_relu_2_2 = h
        h = self.to_relu_3_3(h)
        h_relu_3_3 = h
        h = self.to_relu_4_3(h)
        h_relu_4_3 = h
        h = self.to_relu_5_3(h)
        h_relu_5_3 = h
        out = h_relu_1_2, h_relu_2_2, h_relu_3_3, h_relu_4_3, h_relu_5_3
        return out
